from monai.networks.nets.attentionunet import AttentionUnet
import torch.nn as nn 
import torch 

class AttUNet3D(nn.Module):
    def __init__(self, in_ch, out_ch) -> None:
        super().__init__()
        self.model = AttentionUnet(3, in_ch, out_ch, channels=[64, 128, 256, 512], strides=[2, 2, 2, 2])
    
    def forward(self, x):
        return self.model(x)
    
class AttUNet3D_32(nn.Module):
    def __init__(self, in_ch, out_ch) -> None:
        super().__init__()
        self.model = AttentionUnet(3, in_ch, out_ch, channels=[32, 64, 128, 256], strides=[2, 2, 2, 2])
    
    def forward(self, x):
        return self.model(x)
        


